// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
// SPDX-License-Identifier: MIT

use img_parts::{
    jpeg::{
        markers::{self},
        Jpeg, JpegSegment,
    },
    Bytes,
};
use serde::{Deserialize, Serialize};
use std::{
    fs::{create_dir_all, OpenOptions},
    io::Read,
    path::*,
};
use thiserror::Error;

const MAX_JPEG_MARKER_SIZE: usize = 64000; // 64K
const AIGC_MARKER: [u8; 4] = *b"AIGC";

/// Reads the AIGC metadata from a specified JPEG file.
///
/// # Arguments
///
/// * `asset_path` - A reference to the path of the JPEG file to be read.
///
/// # Returns
///
/// * `Ok(String)` - If successful, returns the AIGC metadata as a JSON string.
/// * `Err(AigcLabelError)` - Returns an error if the file cannot be opened, is not a valid JPEG,
///   or does not contain an AIGC metadata segment.
pub fn read_aigc_meta(asset_path: &Path) -> Result<String> {
    let mut input_stream = OpenOptions::new().read(true).open(asset_path)?;
    let mut jpeg_buf = Vec::new();
    input_stream.read_to_end(&mut jpeg_buf)?;

    let jpeg = Jpeg::from_bytes(jpeg_buf.into())
        .map_err(|err| AigcLabelError::InvalidAsset(format!("Unknown image format: {:?}", err)))?;
    let aigc_seg = get_aigc_segment(&jpeg)?;
    match aigc_seg {
        Some(seg) => Ok(seg.1),
        None => Err(AigcLabelError::ResourceNotFound("AIGC segment".to_string())),
    }
}

/// Saves AIGC metadata to a JPEG file.
///
/// # Arguments
///
/// * `src_path` - The path to the source JPEG file.
/// * `dest_path` - The path where the modified JPEG file will be saved.
/// * `label` - The AIGC label string.
/// * `content_provider` - The name of the content producer or propagator.
/// * `provider_id` - The producerId or propagatorId.
/// * `write_as_producer` - If `true`, writes the metadata as a producer. If `false`, writes as a propagator.
///   Attempting to write as a producer to a file that already has AIGC metadata will result in an error.
///
/// # Returns
///
/// * `Ok(String)` - If successful, returns the full AIGC metadata that was written, as a JSON string.
/// * `Err(AigcLabelError)` - Returns an error for I/O issues, if the source file is invalid,
///   or if there's a logic error (e.g., trying to overwrite producer metadata).
pub fn save_aigc_meta(
    src_path: &std::path::Path,
    dest_path: &std::path::Path,
    label: String,
    content_provider: String,
    provider_id: String,
    write_as_producer: bool,
) -> Result<String> {
    // copy dest file for write.
    if !dest_path.exists() {
        // ensure the path to the file exists
        if let Some(output_dir) = dest_path.parent() {
            create_dir_all(output_dir)?;
        }
        std::fs::copy(&src_path, &dest_path)?;
    }

    let mut input_stream = OpenOptions::new().read(true).open(src_path)?;
    let mut jpeg_buf = Vec::new();
    input_stream.read_to_end(&mut jpeg_buf)?;

    let mut jpeg = Jpeg::from_bytes(jpeg_buf.into())
        .map_err(|err| AigcLabelError::InvalidAsset(format!("Unknown image format: {:?}", err)))?;
    let existing_aigc_segment = get_aigc_segment(&jpeg)?;

    // check if need override exist value.
    let aigc_meta = match existing_aigc_segment {
        Some(seg) => {
            let mut existing_metadata = serde_json::from_slice::<AigcMetadata>(&seg.1.as_bytes())
                .map_err(|_err| {
                AigcLabelError::InvalidAsset("parse exist aigc meta failed".to_string())
            })?;

            if write_as_producer {
                return Err(AigcLabelError::InvalidAsset(format!(
                    "AIGC metadata already exists, cannot overwrite as producer"
                )));
            } else {
                // remove exist segment first.
                jpeg.segments_mut().remove(seg.0);
                // append propagator info.
                existing_metadata.content_propagator = content_provider;
                existing_metadata.propagate_id = provider_id;
                existing_metadata
            }
        }
        None => match write_as_producer {
            true => AigcMetadata {
                label,
                content_producer: content_provider,
                produce_id: provider_id,
                ..Default::default()
            },
            false => AigcMetadata {
                label,
                content_propagator: content_provider,
                propagate_id: provider_id,
                ..Default::default()
            },
        },
    };

    // create AIGC segment.
    let aigc_meta_str = serde_json::to_string(&aigc_meta)
        .map_err(|_err| AigcLabelError::InvalidAsset("to json string failed".to_string()))?;
    let mut final_aigc_segment_data = AIGC_MARKER.to_vec();
    final_aigc_segment_data.extend_from_slice(aigc_meta_str.as_bytes());

    if final_aigc_segment_data.len() > MAX_JPEG_MARKER_SIZE {
        return Err(AigcLabelError::InvalidAsset(format!(
            "jpeg max segment len:{}, now:{}",
            MAX_JPEG_MARKER_SIZE,
            final_aigc_segment_data.len()
        )));
    }

    let segment_bytes = Bytes::from(final_aigc_segment_data);
    let app11_segment = JpegSegment::new_with_contents(markers::APP11, segment_bytes);
    jpeg.segments_mut().insert(1, app11_segment); // we put meta in the beginning.

    let output_stream = OpenOptions::new().read(true).write(true).open(dest_path)?;
    jpeg.encoder()
        .write_to(output_stream)
        .map_err(|_err| AigcLabelError::InvalidAsset("JPEG write error".to_owned()))?;

    Ok(aigc_meta_str)
}

#[derive(Debug, Default, Deserialize, Serialize, Clone)]
struct AigcMetadata {
    #[serde(rename = "Label")]
    pub label: String,

    #[serde(rename = "ContentProducer", default = "default_empty_str")]
    pub content_producer: String,

    #[serde(rename = "ProduceID", default = "default_empty_str")]
    pub produce_id: String,

    #[serde(rename = "ReservedCode1", default = "default_empty_str")]
    pub reserved_code1: String,

    #[serde(rename = "ContentPropagator", default = "default_empty_str")]
    pub content_propagator: String,

    #[serde(rename = "PropagateID", default = "default_empty_str")]
    pub propagate_id: String,

    #[serde(rename = "ReservedCode2", default = "default_empty_str")]
    pub reserved_code2: String,
}
fn default_empty_str() -> String {
    String::new()
}

// first aigc label specific segment
fn get_aigc_segment(jpeg: &img_parts::jpeg::Jpeg) -> Result<Option<(usize, String)>> {
    let segments = jpeg.segments();

    for (i, segment) in segments.iter().enumerate() {
        let raw_bytes = segment.contents();
        let seg_type = segment.marker();

        if raw_bytes.len() > AIGC_MARKER.len() && seg_type == markers::APP11 {
            // we need at least 4 bytes in each segment for AIGC.
            let (marker, data) = raw_bytes.split_at(4);
            if marker == AIGC_MARKER {
                let meta_str = String::from_utf8_lossy(data).to_string();
                return Ok(Some((i, meta_str)));
            }
        }
    }
    Ok(None)
}

/// The Errors that may occur when processing an image.
#[derive(Debug, Error)]
#[repr(i64)]
pub enum AigcLabelError {
    #[error("bad param")]
    BadParam = 1,
    #[error(transparent)]
    IoError(#[from] std::io::Error) = 2,

    #[error("asset could not be parsed: {0}")]
    InvalidAsset(String) = 3,

    #[error("resource not found: {0}")]
    ResourceNotFound(String) = 59,
}

pub fn error_discriminant(v: &AigcLabelError) -> i64 {
    unsafe { *(v as *const AigcLabelError as *const i64) }
}

pub type Result<T> = core::result::Result<T, AigcLabelError>;

#[cfg(test)]
pub mod tests {
    #![allow(clippy::expect_used)]
    #![allow(clippy::panic)]
    #![allow(clippy::unwrap_used)]

    use crate::{read_aigc_meta, save_aigc_meta, AigcLabelError, AigcMetadata};
    use std::{fs, path::PathBuf};

    fn test_file_path(file_name: &str) -> PathBuf {
        let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        path.push("tests");
        path.push(file_name);
        path
    }

    #[test]
    fn test_read_write_jpg() {
        // read without AIGC segment.
        {
            let test_jpeg = test_file_path("test_jpeg.jpeg");
            let result = read_aigc_meta(&test_jpeg);
            assert!(matches!(result, Err(AigcLabelError::ResourceNotFound(_))));
        }

        let write_label = "1";
        let write_producer_id = "1234567890";
        let write_content_producer = "test_producer";
        // insert AIGC segment.
        {
            let src_path = test_file_path("test_jpeg.jpeg");
            let dest_path = test_file_path("with_aigc_temp.jpg");
            if dest_path.exists() {
                fs::remove_file(&dest_path).unwrap_or_else(|_| panic!("rm test file failed"));
            }
            println!("src:{:?} to dest:{:?}", &src_path, &dest_path);

            let result = save_aigc_meta(
                &src_path,
                &dest_path,
                write_label.to_string(),
                write_content_producer.to_string(),
                write_producer_id.to_string(),
                true,
            );
            result.unwrap_or_else(|e| panic!("{:?}", e));
        }

        // read with AIGC seg.
        {
            let with_aigc_file_path = test_file_path("with_aigc_temp.jpg");
            let result = read_aigc_meta(&with_aigc_file_path);
            let result = result.unwrap_or_else(|e| panic!("{:?}", e));
            println!("read aigc meta: {:?}", result);
            let existing_metadata = serde_json::from_slice::<AigcMetadata>(result.as_bytes())
                .unwrap_or_else(|e| panic!("{:?}", e));
            assert_eq!(existing_metadata.label, write_label);
            assert_eq!(existing_metadata.produce_id, write_producer_id);
            assert_eq!(existing_metadata.content_producer, write_content_producer);
            fs::remove_file(&with_aigc_file_path).unwrap();
        }
    }
}
